"""
记忆管理器 - 统一管理多种记忆类型
整合LangChain记忆、Mem0记忆和数据库记忆，提供统一的记忆操作接口
"""

from datetime import datetime
from typing import Dict, Any, List, Optional

from langchain.schema import BaseMemory, HumanMessage, AIMessage
from sqlalchemy.orm import Session

from config.logging_config import get_logger
from pojo.ai_web_memory_his_model import AIWebMemoryHisModel
from .Mem0Interface import Mem0MemoryProvider


class MemoryManager:
    """
    统一记忆管理器
    
    功能：
    1. 管理LangChain记忆（短期对话记忆）
    2. 管理Mem0记忆（智能长期记忆）
    3. 管理数据库记忆（持久化存储）
    4. 提供统一的记忆操作接口
    """

    def __init__(
            self,
            langchain_memory: Optional[BaseMemory] = None,
            mem0_provider: Optional[Mem0MemoryProvider] = None,
            db_session: Session = None,
            scene_id: str = "",
            agent_id: str = "",
            session_id: str = "",
            user_id: str = "1",
            config: Optional[Dict[str, Any]] = None
    ):
        """
        初始化记忆管理器
        
        Args:
            langchain_memory: LangChain记忆实例
            mem0_provider: Mem0记忆提供者
            db_session: 数据库会话
            scene_id: 场景ID
            agent_id: 智能体ID
            session_id: 会话ID
            user_id: 用户ID
            config: 配置信息
        """
        self.langchain_memory = langchain_memory
        self.mem0_provider = mem0_provider
        self.db_session = db_session
        self.scene_id = scene_id
        self.agent_id = agent_id
        self.session_id = session_id
        self.user_id = user_id
        self.config = config or {}

        self.logger = get_logger(__name__)

        # 记忆统计
        self.stats = {
            "messages_processed": 0,
            "db_saves": 0,
            "mem0_saves": 0,
            "last_save_time": None
        }

        self.logger.info(f"记忆管理器初始化: session_id={session_id}, agent_id={agent_id}")

    def add_message(self, human_message: str, ai_message: str) -> bool:
        """
        添加对话消息到所有记忆系统（包括向量库ChromaDB）
        
        Args:
            human_message: 用户消息
            ai_message: AI回复
            
        Returns:
            是否成功
        """
        try:
            success = True

            # 1. 优先添加到数据库记忆（确保持久化存储，支持同步更新）
            if self.db_session and self.config.get("auto_save", True):
                try:
                    db_success = self._save_to_database(human_message, ai_message)
                    if db_success:
                        self.stats["db_saves"] += 1
                        self.logger.debug("数据库记忆保存成功（同步）")
                    else:
                        success = False
                        self.logger.warning("数据库记忆保存失败")
                except Exception as e:
                    self.logger.warning(f"数据库记忆保存失败: {e}")
                    success = False

            # 3. 添加到LangChain记忆（短期记忆，用于快速访问）
            if self.langchain_memory:
                try:
                    self.langchain_memory.save_context(
                        {"input": human_message},
                        {"output": ai_message}
                    )
                    self.logger.debug("LangChain记忆保存成功")
                except Exception as e:
                    self.logger.warning(f"LangChain记忆保存失败: {e}")
                    # LangChain记忆失败不影响整体成功，因为数据库和向量库更重要

            # 更新统计
            self.stats["messages_processed"] += 1
            self.stats["last_save_time"] = datetime.now().isoformat()

            return success

        except Exception as e:
            self.logger.error(f"添加消息失败: {e}")
            return False

    def _save_to_database(self, human_message: str, ai_message: str) -> bool:
        """保存到数据库（同步方式，确保立即持久化）"""
        try:
            if not self.db_session:
                return False

            conversation_content = f"用户: {human_message}\n助手: {ai_message}"

            memory_record = AIWebMemoryHisModel(
                scene_id=self.scene_id,
                agent_id=self.agent_id,
                session_id=self.session_id,
                user_id=self.user_id,
                his_text=conversation_content,
                status=1,  # 临时状态（根据要求）
                importance_level=3,  # 默认重要程度
                create_time=datetime.now()
            )

            # 设置扩展信息
            extra_info = {
                "message_type": "conversation",
                "human_message": human_message,
                "ai_message": ai_message,
                "timestamp": datetime.now().isoformat(),
                "memory_source": "real_time_conversation"  # 标记为实时对话
            }
            memory_record.set_extra_info(extra_info)

            self.db_session.add(memory_record)
            self.db_session.commit()

            self.logger.debug(f"数据库记忆保存成功: session_id={self.session_id}")
            return True

        except Exception as e:
            self.logger.error(f"保存到数据库失败: {e}")
            if self.db_session:
                self.db_session.rollback()
            return False

    def get_conversation_history(self, limit: int = 10) -> List[Dict[str, Any]]:
        """
        获取对话历史
        
        Args:
            limit: 返回的消息数量限制
            
        Returns:
            对话历史列表
        """
        try:
            history = []

            # 优先从LangChain记忆获取（最新的对话）
            if self.langchain_memory and hasattr(self.langchain_memory, 'chat_memory'):
                messages = self.langchain_memory.chat_memory.messages
                for message in messages[-limit * 2:]:  # 获取最近的消息，*2因为有用户和AI消息
                    if isinstance(message, HumanMessage):
                        history.append({
                            "role": "user",
                            "content": message.content,
                            "timestamp": datetime.now().isoformat(),
                            "source": "langchain"
                        })
                    elif isinstance(message, AIMessage):
                        history.append({
                            "role": "assistant",
                            "content": message.content,
                            "timestamp": datetime.now().isoformat(),
                            "source": "langchain"
                        })

            # 如果LangChain记忆中的消息不够，从数据库补充
            if len(history) < limit and self.db_session:
                remaining_limit = limit - len(history)

                db_memories = self.db_session.query(AIWebMemoryHisModel).filter(
                    AIWebMemoryHisModel.session_id == self.session_id
                ).order_by(AIWebMemoryHisModel.create_time.desc()).limit(remaining_limit).all()

                for memory in reversed(db_memories):  # 按时间正序
                    # 尝试解析对话内容
                    extra_info = memory.get_extra_info()
                    if extra_info and "human_message" in extra_info and "ai_message" in extra_info:
                        history.extend([
                            {
                                "role": "user",
                                "content": extra_info["human_message"],
                                "timestamp": memory.create_time.isoformat(),
                                "source": "database"
                            },
                            {
                                "role": "assistant",
                                "content": extra_info["ai_message"],
                                "timestamp": memory.create_time.isoformat(),
                                "source": "database"
                            }
                        ])
                    else:
                        # 如果没有结构化数据，解析原始文本
                        content = memory.his_text
                        if "用户:" in content and "助手:" in content:
                            parts = content.split("助手:")
                            if len(parts) == 2:
                                user_part = parts[0].replace("用户:", "").strip()
                                ai_part = parts[1].strip()
                                history.extend([
                                    {
                                        "role": "user",
                                        "content": user_part,
                                        "timestamp": memory.create_time.isoformat(),
                                        "source": "database"
                                    },
                                    {
                                        "role": "assistant",
                                        "content": ai_part,
                                        "timestamp": memory.create_time.isoformat(),
                                        "source": "database"
                                    }
                                ])

            return history[-limit:] if len(history) > limit else history

        except Exception as e:
            self.logger.error(f"获取对话历史失败: {e}")
            return []

    def search_relevant_memories(
            self,
            query: str,
            limit: int = 5
    ) -> List[Dict[str, Any]]:
        """
        搜索相关记忆（同步方法）
        
        Args:
            query: 搜索查询
            limit: 返回数量限制
            
        Returns:
            相关记忆列表
        """
        try:
            memories = []

            # 从Mem0搜索智能记忆（同步调用）
            if self.mem0_provider:
                # 直接调用Mem0的同步搜索方法
                mem0_search_results = self.mem0_provider.mem0.search_memories(
                    query=query,
                    user_id=self.user_id,
                    agent_id=self.agent_id,
                    scene_id=self.scene_id,
                    limit=limit
                )

                # 处理搜索结果
                for result in mem0_search_results:
                    content = result.get("memory", result.get("content", ""))
                    memories.append({
                        "content": content,
                        "source": "mem0",
                        "relevance": "high",
                        "timestamp": datetime.now().isoformat(),
                        "score": result.get("score", 0.0),
                        "metadata": result.get("metadata", {})
                    })

            # 如果Mem0结果不足，从数据库搜索
            if len(memories) < limit and self.db_session:
                remaining_limit = limit - len(memories)

                db_memories = self.db_session.query(AIWebMemoryHisModel).filter(
                    AIWebMemoryHisModel.scene_id == self.scene_id,
                    AIWebMemoryHisModel.agent_id == self.agent_id,
                    AIWebMemoryHisModel.user_id == self.user_id,
                    AIWebMemoryHisModel.his_text.like(f"%{query}%")
                ).order_by(
                    AIWebMemoryHisModel.importance_level.desc(),
                    AIWebMemoryHisModel.create_time.desc()
                ).limit(remaining_limit).all()

                for memory in db_memories:
                    memories.append({
                        "content": memory.his_text,
                        "source": "database",
                        "relevance": "medium",
                        "timestamp": memory.create_time.isoformat(),
                        "importance": memory.importance_level
                    })

            return memories

        except Exception as e:
            self.logger.error(f"搜索相关记忆失败: {e}")
            return []

    def preload_session_history(self, limit: int = 20) -> bool:
        """
        预加载会话历史到LangChain记忆中
        
        Args:
            limit: 加载的历史消息数量限制
            
        Returns:
            是否成功
        """
        try:
            if not self.langchain_memory or not self.db_session:
                return False

            # 从数据库获取历史对话
            db_memories = self.db_session.query(AIWebMemoryHisModel).filter(
                AIWebMemoryHisModel.session_id == self.session_id
            ).order_by(AIWebMemoryHisModel.create_time.asc()).limit(limit).all()

            loaded_count = 0
            for memory in db_memories:
                extra_info = memory.get_extra_info()
                if extra_info and "human_message" in extra_info and "ai_message" in extra_info:
                    try:
                        self.langchain_memory.save_context(
                            {"input": extra_info["human_message"]},
                            {"output": extra_info["ai_message"]}
                        )
                        loaded_count += 1
                    except Exception as e:
                        self.logger.warning(f"预加载消息失败: {e}")
                        continue

            self.logger.info(f"预加载完成: 加载了{loaded_count}条历史消息")
            return loaded_count > 0

        except Exception as e:
            self.logger.error(f"预加载会话历史失败: {e}")
            return False

    def clear_memory(self, clear_db: bool = False, clear_mem0: bool = False) -> bool:
        """
        清空记忆
        
        Args:
            clear_db: 是否清空数据库记忆
            clear_mem0: 是否清空Mem0记忆
            
        Returns:
            是否成功
        """
        try:
            success = True

            # 清空LangChain记忆
            if self.langchain_memory:
                try:
                    self.langchain_memory.clear()
                    self.logger.info("LangChain记忆已清空")
                except Exception as e:
                    self.logger.warning(f"清空LangChain记忆失败: {e}")
                    success = False

            # 清空数据库记忆（可选）
            if clear_db and self.db_session:
                try:
                    deleted_count = self.db_session.query(AIWebMemoryHisModel).filter(
                        AIWebMemoryHisModel.session_id == self.session_id
                    ).delete()
                    self.db_session.commit()
                    self.logger.info(f"数据库记忆已清空: 删除了{deleted_count}条记录")
                except Exception as e:
                    self.logger.warning(f"清空数据库记忆失败: {e}")
                    self.db_session.rollback()
                    success = False

            # 清空Mem0记忆（实现真正的删除功能）
            if clear_mem0 and self.mem0_provider:
                try:
                    # 尝试异步删除Mem0记忆
                    import asyncio
                    try:
                        # 尝试在现有事件循环中运行
                        loop = asyncio.get_event_loop()
                        if loop.is_running():
                            # 在运行中的事件循环中，创建任务
                            asyncio.create_task(self._clear_mem0_memories())
                            self.logger.info("Mem0记忆清空任务已创建（异步）")
                        else:
                            # 事件循环未运行，同步执行
                            mem0_success = loop.run_until_complete(self._clear_mem0_memories())
                            if mem0_success:
                                self.logger.info("Mem0记忆已清空（同步）")
                            else:
                                self.logger.warning("Mem0记忆清空失败")
                                success = False
                    except RuntimeError:
                        # 没有事件循环，创建新的
                        mem0_success = asyncio.run(self._clear_mem0_memories())
                        if mem0_success:
                            self.logger.info("Mem0记忆已清空（新事件循环）")
                        else:
                            self.logger.warning("Mem0记忆清空失败")
                            success = False
                except Exception as e:
                    self.logger.warning(f"清空Mem0记忆失败: {e}")
                    success = False

            return success

        except Exception as e:
            self.logger.error(f"清空记忆失败: {e}")
            return False

    async def _clear_mem0_memories(self) -> bool:
        """
        清空Mem0记忆的异步实现
        
        Returns:
            是否成功
        """
        try:
            if not self.mem0_provider:
                return False

            # 根据Mem0的API文档实现删除功能
            # 方法1: 使用delete方法删除特定会话的记忆
            try:
                # 如果Mem0提供按metadata删除的功能
                delete_success = await self.mem0_provider.delete_memories_by_metadata(
                    user_id=self.user_id,
                    agent_id=self.agent_id,
                    scene_id=self.scene_id,
                    session_id=self.session_id
                )

                if delete_success:
                    self.logger.info(f"✅ 成功删除Mem0记忆: session_id={self.session_id}")
                    return True

            except AttributeError:
                # 如果不支持按metadata删除，尝试其他方法
                self.logger.debug("Mem0不支持按metadata删除，尝试其他方法")

            # 方法2: 搜索相关记忆然后逐个删除
            try:
                # 获取所有相关记忆的ID
                memory_ids = await self.mem0_provider.get_memory_ids_by_session(
                    user_id=self.user_id,
                    agent_id=self.agent_id,
                    scene_id=self.scene_id,
                    session_id=self.session_id
                )

                if memory_ids:
                    deleted_count = 0
                    for memory_id in memory_ids:
                        try:
                            if await self.mem0_provider.delete_memory_by_id(memory_id):
                                deleted_count += 1
                        except Exception as delete_error:
                            self.logger.warning(f"删除Mem0记忆 {memory_id} 失败: {delete_error}")

                    self.logger.info(f"✅ 成功删除 {deleted_count}/{len(memory_ids)} 条Mem0记忆")
                    return deleted_count > 0
                else:
                    self.logger.info("没有找到要删除的Mem0记忆")
                    return True

            except AttributeError:
                # 如果也不支持按ID删除，尝试基础方法
                self.logger.debug("Mem0不支持按ID删除，尝试基础方法")

            # 方法3: 使用基础的delete方法（如果Mem0实例有delete方法）
            try:
                if hasattr(self.mem0_provider.mem0, 'mem0_instance'):
                    mem0_instance = self.mem0_provider.mem0.mem0_instance

                    # 尝试使用Mem0的delete方法
                    if hasattr(mem0_instance, 'delete'):
                        # 构建删除查询条件
                        delete_filters = {
                            "user_id": self.user_id,
                            "agent_id": self.agent_id,
                            "scene_id": self.scene_id,
                            "session_id": self.session_id
                        }

                        delete_result = mem0_instance.delete(filters=delete_filters)
                        self.logger.info(f"✅ 使用基础delete方法删除Mem0记忆: {delete_result}")
                        return True
                    elif hasattr(mem0_instance, 'delete_all'):
                        # 如果有delete_all方法，谨慎使用
                        self.logger.warning("⚠️ 只能使用delete_all方法，这会删除所有记忆，跳过")
                        return False

            except Exception as basic_error:
                self.logger.warning(f"基础删除方法失败: {basic_error}")

            # 如果所有方法都失败，记录警告但不失败
            self.logger.warning("⚠️ 无法删除Mem0记忆，可能需要手动清理或等待过期")
            return True  # 返回True因为这不应该阻止其他清理操作

        except Exception as e:
            self.logger.error(f"清空Mem0记忆异步实现失败: {e}")
            return False

    def get_memory_stats(self) -> Dict[str, Any]:
        """获取记忆统计信息"""
        stats = self.stats.copy()

        # 添加LangChain记忆统计
        if self.langchain_memory and hasattr(self.langchain_memory, 'chat_memory'):
            stats["langchain_messages"] = len(self.langchain_memory.chat_memory.messages)
        else:
            stats["langchain_messages"] = 0

        # 添加数据库记忆统计
        if self.db_session:
            try:
                db_count = self.db_session.query(AIWebMemoryHisModel).filter(
                    AIWebMemoryHisModel.session_id == self.session_id
                ).count()
                stats["database_messages"] = db_count
            except:
                stats["database_messages"] = 0
        else:
            stats["database_messages"] = 0

        # 添加配置信息
        stats["config"] = self.config
        stats["memory_types"] = []

        if self.langchain_memory:
            stats["memory_types"].append("langchain")
        if self.mem0_provider:
            stats["memory_types"].append("mem0")
        if self.db_session:
            stats["memory_types"].append("database")

        return stats

    def get_langchain_memory(self) -> Optional[BaseMemory]:
        """获取LangChain记忆实例（用于与Agent集成）"""
        return self.langchain_memory

    def set_importance_level(self, memory_id: int, importance: int) -> bool:
        """
        设置记忆重要程度
        
        Args:
            memory_id: 记忆ID
            importance: 重要程度（1-5）
            
        Returns:
            是否成功
        """
        try:
            if not self.db_session:
                return False

            memory = self.db_session.query(AIWebMemoryHisModel).filter(
                AIWebMemoryHisModel.id == memory_id
            ).first()

            if memory:
                memory.importance_level = importance
                memory.update_time = datetime.now()
                self.db_session.commit()
                return True

            return False

        except Exception as e:
            self.logger.error(f"设置记忆重要程度失败: {e}")
            if self.db_session:
                self.db_session.rollback()
            return False

    async def backup_memory(self) -> Dict[str, Any]:
        """
        备份当前会话的记忆数据
        
        Returns:
            备份数据
        """
        try:
            backup_data = {
                "session_id": self.session_id,
                "agent_id": self.agent_id,
                "scene_id": self.scene_id,
                "user_id": self.user_id,
                "backup_time": datetime.now().isoformat(),
                "langchain_messages": [],
                "database_memories": []
            }

            # 备份LangChain记忆
            if self.langchain_memory and hasattr(self.langchain_memory, 'chat_memory'):
                for message in self.langchain_memory.chat_memory.messages:
                    backup_data["langchain_messages"].append({
                        "type": type(message).__name__,
                        "content": message.content,
                        "timestamp": datetime.now().isoformat()
                    })

            # 备份数据库记忆
            if self.db_session:
                db_memories = self.db_session.query(AIWebMemoryHisModel).filter(
                    AIWebMemoryHisModel.session_id == self.session_id
                ).all()

                for memory in db_memories:
                    backup_data["database_memories"].append({
                        "id": memory.id,
                        "content": memory.his_text,
                        "create_time": memory.create_time.isoformat(),
                        "importance_level": memory.importance_level,
                        "status": memory.status,
                        "extra_info": memory.get_extra_info()
                    })

            return backup_data

        except Exception as e:
            self.logger.error(f"备份记忆失败: {e}")
            return {}

    def __repr__(self) -> str:
        """字符串表示"""
        types = []
        if self.langchain_memory:
            types.append("LangChain")
        if self.mem0_provider:
            types.append("Mem0")
        if self.db_session:
            types.append("Database")

        return f"MemoryManager(session_id='{self.session_id}', types={types})"
